#include "StdAfx.h"
#include "ProcessList.h"

CProcessList::CProcessList(void)
{
}

CProcessList::~CProcessList(void)
{
}

void CProcessList::Destroy() {
	for(POSITION pos = m_Processes.GetHeadPosition(); pos; )
		delete m_Processes.GetNext(pos);
	m_Processes.RemoveAll();
}

const CProcessList::ProcessList& CProcessList::Create() {
	Destroy();

	HANDLE hSnapshot = ::CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
	PROCESSENTRY32 pe = { sizeof(pe) };
	if(::Process32First(hSnapshot, &pe)) {
		while(::Process32Next(hSnapshot, &pe)) {
			HANDLE h = ::OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, pe.th32ProcessID);
			if(h) {
#ifndef _WIN64
				BOOL iswow = FALSE;
				::IsWow64Process(h, &iswow);
				if(iswow) {
#endif
				m_Processes.AddTail(new PROCESSENTRY32(pe));
				::CloseHandle(h);
#ifndef _WIN64
				}
#endif
			}
		}
	}
	::CloseHandle(hSnapshot);
	return m_Processes;
}

bool CProcessList::IsSystem(DWORD pid) {
	HANDLE hProcess = ::OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, pid);	
	if(hProcess == NULL) return true;

	HANDLE hToken = NULL;
	::OpenProcessToken(hProcess, TOKEN_QUERY, &hToken);
	if(hToken == NULL) {
		::CloseHandle(hProcess);
		return true;
	}

	TOKEN_USER *user = (TOKEN_USER*)::malloc(128);
	DWORD size;
	BOOL ok = ::GetTokenInformation(hToken, TokenUser, user, 128, &size);
	ASSERT(ok);
	PSID sid = user->User.Sid;
	BOOL system = ::IsWellKnownSid(sid, WinLocalSystemSid) | ::IsWellKnownSid(sid, WinNetworkServiceSid) | ::IsWellKnownSid(sid, WinLocalServiceSid);
	::CloseHandle(hToken);
	::CloseHandle(hProcess);
	::free(user);

	return system ? true : false;
}

bool CProcessList::IsWow64Process(DWORD pid) {
	HANDLE hProcess = ::OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, pid);
	if(!hProcess) return false;

	BOOL iswow;
	::IsWow64Process(hProcess, &iswow);
	::CloseHandle(hProcess);
	return iswow ? true : false;
}

// CModuleList

CModuleList::CModuleList(DWORD pid) {
	m_hEnum = ::CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, pid);
	Reset();
}

void CModuleList::Reset() {
	::ZeroMemory(&m_ModuleInfo, sizeof(m_ModuleInfo));
	m_ModuleInfo.dwSize = sizeof(m_ModuleInfo);
}

CModuleList::~CModuleList() {
	::CloseHandle(m_hEnum);
}

MODULEENTRY32* CModuleList::GetNext() {
	BOOL ok;
	if(m_ModuleInfo.th32ProcessID == 0)
		ok = ::Module32First(m_hEnum, &m_ModuleInfo);
	else
		ok = ::Module32Next(m_hEnum, &m_ModuleInfo);
	return ok ? &m_ModuleInfo : NULL;
}

// CThreadList

CThreadList::CThreadList() {
	m_hEnum = ::CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
	Reset();
}

void CThreadList::Reset() {
	::ZeroMemory(&m_ThreadInfo, sizeof(m_ThreadInfo));
	m_ThreadInfo.dwSize = sizeof(m_ThreadInfo);
	m_First = true;
}

CThreadList::~CThreadList() {
	::CloseHandle(m_hEnum);
}

THREADENTRY32* CThreadList::GetNext() {
	BOOL ok;
	if(m_First) {
		ok = ::Thread32First(m_hEnum, &m_ThreadInfo);
		m_First = false;
	}
	else
		ok = ::Thread32Next(m_hEnum, &m_ThreadInfo);
	return ok ? &m_ThreadInfo : NULL;
}

// CHeapList

CHeapList::CHeapList(DWORD pid) {
	m_hEnum = ::CreateToolhelp32Snapshot(TH32CS_SNAPHEAPLIST, pid);
	::memset(&m_Heap, 0, sizeof(m_Heap));
	m_Heap.dwSize = sizeof(m_Heap);
}

CHeapList::~CHeapList() {
	::CloseHandle(m_hEnum);
}

HEAPLIST32* CHeapList::GetNext() {
	BOOL more;
	if(m_Heap.th32ProcessID == 0)
		more = ::Heap32ListFirst(m_hEnum, &m_Heap);
	else
		more = ::Heap32ListNext(m_hEnum, &m_Heap);
	return more ? &m_Heap : NULL;
}
